import sys
import numpy as np
import pandas as pd
from sklearn.cluster import KMeans
from scipy.stats import multivariate_normal as mvn

from fig_utils import *

"""

    GMM ( Gaussian Mixture Model ) Implementation

"""


class GMM(object):
    def __init__(self, data, num_clusters, max_iter=50, interval=0.1):
        self._data = data
        self._N = self._data.shape[0]
        self._k = num_clusters
        self._max_iter = max_iter
        self._log = np.zeros((self._k, self._N))
        self._interval = interval
        self._poster = None
        self._mu = None
        self._cov = None
        self._prior = None

    def fit(self, init='kmeans', visual=True):
        """
        Estimate GMM parameters
        :return: None
        """
        if init == 'kmeans':
            # Initialize GMM parameters with k-means
            self._init_kmeans()
        elif init == 'random':
            # Randomly initial GMM parameters
            self._init_random()
        else:
            raise ValueError('init only assigned from [kmeans, random]')

        log_likelihoods = []
        if visual:
            result_fig(self._data)

        for i in range(self._max_iter + 1):
            self._e_step()
            self._m_step()
            predict = self.predict()
            log_likelihoods.append(self.log_likelihood())
            if visual:
                result_fig(self._data, log_likelihoods, i, self._max_iter,
                           self._k, predict, self._mu, self._interval)

    def predict(self):
        """
        Predict labels
        :return:  labels of data points
        """
        res = np.argmax(self._poster, axis=0)
        return res

    def log_likelihood(self):
        for k in range(self._k):
            self._log[k] = mvn.pdf(self._data,
                                   mean=self._mu[k],
                                   cov=self._cov[k])
        self._log = np.dot(
            np.diag(self._prior.ravel()), self._log
        )
        ret = self._log.sum(axis=0)

        return np.log(ret).sum()

    @property
    def get_mu(self):
        return self._mu

    @property
    def get_var(self):
        return self._cov

    def _init_kmeans(self):
        """
        Set initial GMM params with k-means initialization
        :return: None
        """
        k_means = KMeans(init='k-means++', n_clusters=self._k)
        k_means.fit(self._data)
        category = k_means.labels_

        self._poster = np.zeros((self._k, self._N))
        self._mu = k_means.cluster_centers_
        self._cov = np.asarray(
            [np.cov(self._data[category == k], rowvar=False) for k in range(self._k)]
        )

        # compute count of different categories
        value_counts = pd.Series(category).value_counts()

        self._prior = np.asarray(
            [value_counts[k] / self._N for k in range(self._k)]
        ).reshape((self._k, 1))

    def _init_random(self):
        """
        Randomly initialize GMM parameters
        :return: None
        """
        self._poster = np.zeros((self._k, self._N))

        self._mu = self._data[np.random.choice(self._N, size=self._k, replace=False)]
        self._cov = np.asarray([np.cov(self._data, rowvar=False)] * self._k)
        self._prior = np.ones((self._k, 1)) / self._k

    def _e_step(self):
        """
        Perform E-step
        Update poster responsibility
        :return: None
        """
        for k in range(self._k):
            self._poster[k] = mvn.pdf(self._data,
                                      mean=self._mu[k],
                                      cov=self._cov[k])
        self._poster = np.dot(
            np.diag(self._prior.ravel()), self._poster
        )
        self._poster /= np.sum(self._poster, axis=0)

    def _m_step(self):
        """
        Performs M-step of the GMM
        update priori, mu, covariance matrix
        :return:
        """
        effective_count = np.sum(self._poster, axis=1)

        self._mu = np.asarray(
            [np.dot(self._poster[k], self._data) / effective_count[k]
             for k in range(self._k)]
        )

        self._cov = np.asarray(
            [np.dot((self._data - self._mu[k]).T,
                    np.dot(np.diag(self._poster[k].ravel()),
                           self._data - self._mu[k])
                    ) / effective_count[k] for k in range(self._k)])

        self._prior = (effective_count / self._N).reshape((self._k, 1))
        return


if __name__ == '__main__':
    pass
